package com.sjw.workflow.core;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.List;
import java.util.ArrayList;

/**
 * Workflow上下文，用于在Agent之间传递数据
 */
public class WorkflowContext {
    
    private final Map<String, Object> globalData;
    private final Map<String, AgentResult> agentResults;
    private final List<String> executionLog;
    private final String workflowId;
    private final long startTime;
    
    public WorkflowContext(String workflowId) {
        this.workflowId = workflowId;
        this.globalData = new ConcurrentHashMap<>();
        this.agentResults = new ConcurrentHashMap<>();
        this.executionLog = new ArrayList<>();
        this.startTime = System.currentTimeMillis();
    }
    
    /**
     * 设置全局数据
     */
    public void setGlobalData(String key, Object value) {
        globalData.put(key, value);
    }
    
    /**
     * 获取全局数据
     */
    public Object getGlobalData(String key) {
        return globalData.get(key);
    }
    
    /**
     * 获取所有全局数据
     */
    public Map<String, Object> getAllGlobalData() {
        return new ConcurrentHashMap<>(globalData);
    }
    
    /**
     * 设置Agent执行结果
     */
    public void setAgentResult(String agentId, AgentResult result) {
        agentResults.put(agentId, result);
        executionLog.add(String.format("[%d] Agent %s completed with status: %s", 
            System.currentTimeMillis() - startTime, agentId, result.getStatus()));
    }
    
    /**
     * 获取Agent执行结果
     */
    public AgentResult getAgentResult(String agentId) {
        return agentResults.get(agentId);
    }
    
    /**
     * 检查Agent是否已完成
     */
    public boolean isAgentCompleted(String agentId) {
        AgentResult result = agentResults.get(agentId);
        return result != null && (result.getStatus() == ExecutionStatus.SUCCESS || 
                                 result.getStatus() == ExecutionStatus.FAILED);
    }
    
    /**
     * 检查Agent是否成功完成
     */
    public boolean isAgentSuccessful(String agentId) {
        AgentResult result = agentResults.get(agentId);
        return result != null && result.getStatus() == ExecutionStatus.SUCCESS;
    }
    
    /**
     * 获取所有Agent结果
     */
    public Map<String, AgentResult> getAllAgentResults() {
        return new ConcurrentHashMap<>(agentResults);
    }
    
    /**
     * 添加执行日志
     */
    public void addLog(String message) {
        executionLog.add(String.format("[%d] %s", System.currentTimeMillis() - startTime, message));
    }
    
    /**
     * 获取执行日志
     */
    public List<String> getExecutionLog() {
        return new ArrayList<>(executionLog);
    }
    
    /**
     * 获取Workflow ID
     */
    public String getWorkflowId() {
        return workflowId;
    }
    
    /**
     * 获取执行时间
     */
    public long getExecutionTime() {
        return System.currentTimeMillis() - startTime;
    }
    
    /**
     * 合并其他上下文的数据
     */
    public void mergeContext(WorkflowContext other) {
        if (other != null) {
            globalData.putAll(other.globalData);
            agentResults.putAll(other.agentResults);
            executionLog.addAll(other.executionLog);
        }
    }
}
